from __future__ import print_function, division
import time
#Importing relevant modules and packages
import numpy as np
import matplotlib.pyplot as plt
import random
from RichardsModel_1D import *
from Parameters import *
from scipy import linalg, integrate, io
from Prejacobian_disc import F1,F2,Fy
from tqdm import tqdm



p=Loam() #Soil parameters

#Space and time parameters
Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
DeltaT,Nsim=time_parameters()


# u=np.zeros(3*120)
# uu=np.zeros(Nsim)
# for i in range(3*120):
#     if i < 12:
#         u[i]=-1e-06 #Moisture is supplied only at the first time instant
#     else:
#         u[i]=0
# uu=np.tile(u,(1,120))
# uu=np.transpose(uu)

u=np.zeros(6*120)
uu=np.zeros(Nsim)
for i in range(6*120):
    if i < 26:
        u[i]=-1e-06 #Moisture is supplied only at the first time instant
    else:
        u[i]=0
uu=np.tile(u,(1,120))
uu=np.transpose(uu)
N_obs=3
Kc=0.88*np.ones(Nsim)      #Crop coefficient
# Kc[2700:5400]=0.98
# Kc[5400:8100]=0.88
Kc[5040:Nsim]=1.08
Kc1=1.8*np.ones(Nsim)

Et=1.389e-08*np.ones(Nsim) #Reference evapotranspiration
# Et[2700:5400]=1.489e-08
# Et[5400:8100]=1.389e-08
Et[5040:Nsim]=1.489e-08
Et1=1.289e-08*np.ones(Nsim) #Reference evapotranspiration

Sigma_p=0.01
Sigma_w=0.1
Sigma_v=150
P_pre=np.zeros((Nsim+1,Nz,Nz))
P_cur=np.zeros((Nsim+1,Nz,Nz))
K=np.zeros((Nsim+1,Nz,N_obs))
S_k=np.zeros((Nsim+1,N_obs,N_obs))

# These have been predefined because the dimensions of the states are fixed
#P_pre=100*np.ones((Nx,Nx))
P_cur[0,:,:]=np.diag(Sigma_p*np.ones((Nz,)))
Q=np.diag(Sigma_w*np.ones((Nz,)))
R=np.diag(Sigma_v*np.ones((N_obs,)))




#Initial condition [pressure head]
x0=-0.8*np.ones(Nz)


#Numpy array to store the simulation results [pressure head]
xx=np.zeros((Nsim+1,Nz)) #true state
xx1=np.zeros((Nsim+1,Nz)) #open loop
xx_cur=np.zeros((Nsim+1,Nz)) #filter
xx_pre = np.zeros((Nsim+1,Nz)) #predict
xx[0,:]=x0
xx1[0,:]=1.1*x0
xx_cur[0,:]=1.1*x0

# array for filtering
x_out = np.zeros((Nsim+1,Nz))
x_est = np.zeros((Nsim+1,Nz))

x_est[0,:] = xx_cur[0,:] # time by time output of the particle filters estimate


# setting noise
www1 = 4e-9 # process noise
# vv1 = np.array([8e-7,8e-7]) # measurement noise
vv1=8e-7*np.ones(N_obs)
# vv1 = 8e-07*np.ones([1,16])
# vv1 = np.array([1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6]) # measurement noise
vv2 = np.diag(vv1)
vvv1 = 4e-9 # particle noise
wwww1 = 4e-9

# www1 = 0 # process noise
# # vv1 = np.array([8e-7,8e-7]) # measurement noise
# vv1=0*np.ones(N_obs)
# # vv1 = 8e-07*np.ones([1,16])
# # vv1 = np.array([1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6]) # measurement noise
# vv2 = np.diag(vv1)
# vvv1 = 0 # particle noise
# wwww1 = 0


np.random.seed(1)
ww=0*np.random.normal(0,np.sqrt(0.01),[Nsim,Nz]) #If no model uncertainty is present, ww is multiplied by 0
# np.random.seed(2)
# www=0.000001*np.random.randn(Nsim,Nz)
www = np.random.normal(0,np.sqrt(www1),[Nsim+1,Nz])
# np.random.seed(3)
# vv = np.random.normal(0,np.sqrt(vv1),[Nsim+1,Nz]) #measurement noise
vv = np.random.randn(Nsim+1,N_obs)*vv1


# np.random.seed(5)
wwww = np.random.normal(0,np.sqrt(wwww1),[Nsim,Nz])


#Numpy array to store the simulation results [volumetric moisture content]
yy=np.zeros((Nsim+1,N_obs))
yy1=np.zeros((Nsim+1,N_obs))
yy[0,:]=getOutputs1D(x0).full().ravel()+vv[0,:]
yy1[0,:]=getOutputs1D(xx1[0,:]).full().ravel()+vv[0,:]

# C=np.hstack((np.zeros((N_obs,Nz-N_obs)),np.eye(N_obs),np.zeros((N_obs,Nz))))
# C1=np.hstack((np.eye(Nz),np.zeros((Nz,Nz))))
# C=np.delete(C1,(0,3,5,7,9,11,13),axis=0)

#unknown inputs
M=np.ones((Nz,1))
a1=0*np.ones((Nsim+1,Nz))
a=0.00000*np.ones((Nsim+1,Nz))
# a[:,0:7]=0.00003
# a[:,8:16]=0.00002
# for i in range (Nsim+1):
#     a[i,0:7]=a[i,0:7]+(1e-06)*np.cos(0.0005*i)
#     a[i,8:16]=a[i,8:16]+(1e-06)*np.sin(0.0005*i)
a_est=0.00000*np.ones((Nsim+1,Nz))
# a_est[0,0:7]=0.00001
# a_est[0,8:16]=0.00000
a2=0.00000*np.ones((Nsim+1,Nz))
# a_est[:,0:7]=0.00001
# a_est[:,8:16]=0.00000



# #unknown inputs
# a1=0*np.ones((Nsim+1,Nz))
# a=0.00005*np.ones((Nsim+1,Nz))
# a_est=0.00004*np.ones((Nsim+1,Nz))

# gamma=0.00025
gamma=0.45*np.ones(Nz)*[0.00015,0.00015,0.00015,0.00015,0.00017,0.00017,0.00017,0.00017,0.00016,0.00015,0.00017,0.00017,0.00018,0.00020,0.00022,0.00027]
# gamma=0.0001*np.ones(Nz)*[0.0025,0.0025,0.0020,0.0020,0.0020,0.0020,0.0020,0.0020,0.0020,0.0020,0.0022,0.0025,0.0027,0.0030,0.0035,0.0040]
# gamma=[0.000008,0.000008,0.000008,0.000008,0.00006,0.00006,0.000062,0.000064,0.000066,0.000070,0.000076,0.00008,0.000092,0.000098,0.000104,0.000110]


e0 = time.time()
# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE

    
for i in tqdm(range(1, Nsim+1)):
    # Generate observed data (x bias)
    xx[i,:]= DM(ODE_discrete(tuple(xx[i-1]), uu[i-1],ww[i-1],Kc[i-1],Et[i-1],DeltaT)).full().ravel()+a[i-1,:]+www[i-1]    
    yy[i,:]=getOutputs1D(xx[i,:]).full().ravel()+vv[i,:]
    
np.savetxt('xx.txt',xx)

for i in tqdm(range(1, Nsim+1)):
    # open loop
    xx1[i,:]= DM(ODE_discrete(tuple(xx1[i-1]), uu[i-1],ww[i-1],Kc1[i-1],Et1[i-1],DeltaT)).full().ravel()+a2[i-1,:]+www[i-1]
    
    yy1[i,:]=getOutputs1D(xx1[i,:]).full().ravel()+vv[i,:]
    
    
# REM-EKF Algorithm
# for i in tqdm(range(1, Nsim+1)):
#     xx_pre[i,:] = ODE_discrete(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc[i-1],Et[i-1],DeltaT)+a_est[i-1,:]
    
#     P_pre[i,:,:] = mtimes(mtimes(F1(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc[i-1],Et[i-1],DeltaT),P_cur[i-1,:,:]),np.transpose(F1(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc[i-1],Et[i-1],DeltaT)))+Q
    
#     S_k = mtimes(mtimes(mtimes(C,Fy(tuple(xx_pre[i]))),P_pre[i,:,:]),mtimes(np.transpose(Fy(tuple(xx_pre[i]))),np.transpose(C)))+R
    
#     K[i,:,:] = mtimes(mtimes(P_pre[i,:,:],mtimes(np.transpose(Fy(tuple(xx_pre[i]))),np.transpose(C))),np.linalg.inv(S_k))
    
#     xx_cur[i,:] = xx_pre[i,:]+np.transpose(mtimes(K[i,:,:],yy[i,:]-getOutputs1D(xx_pre[i,:]).full().ravel()))
    
#     P_cur[i,:,:] = mtimes((np.eye(Nz)-mtimes(mtimes(K[i,:,:],C),Fy(tuple(xx_pre[i])))),P_pre[i,:,:])
    
#     a_est[i,:]=(np.ones(Nz)-gamma)*a_est[i-1,:]+gamma*(xx_cur[i,:]-ODE_discrete(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc[i-1],Et[i-1],DeltaT))

for i in tqdm(range(1, Nsim+1)):
    xx_pre[i,:] = DM(ODE_discrete(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc1[i-1],Et1[i-1],DeltaT)).full().ravel()+a_est[i-1,:]
    
    P_pre[i,:,:] = mtimes(mtimes(F1(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc1[i-1],Et1[i-1],DeltaT),P_cur[i-1,:,:]),np.transpose(F1(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc1[i-1],Et1[i-1],DeltaT)))+Q
    
    S_k[i,:,:] = mtimes(mtimes(Fy(tuple(xx_pre[i])),P_pre[i,:,:]),np.transpose(Fy(tuple(xx_pre[i]))))+R
    
    K[i,:,:] = mtimes(mtimes(P_pre[i,:,:],np.transpose(Fy(tuple(xx_pre[i])))),np.linalg.inv(S_k[i,:,:]))
    
    xx_cur[i,:] = xx_pre[i,:]+np.transpose(mtimes(K[i,:,:],yy[i,:]-getOutputs1D(xx_pre[i,:]).full().ravel()))
    
    P_cur[i,:,:] = mtimes((np.eye(Nz)-mtimes(K[i,:,:],Fy(tuple(xx_pre[i])))),P_pre[i,:,:])
    
    a_est[i,:]=(np.ones(Nz)-gamma)*a_est[i-1,:]+gamma*(xx_cur[i,:]-DM(ODE_discrete(tuple(xx_cur[i-1]), uu[i-1],ww[i-1],Kc1[i-1],Et1[i-1],DeltaT))).full().ravel()
    

    

t= list(range(Nsim+1))
# plt.figure(1)
# for i in range (Nz):
#     plt.subplot(5, 7, i+1)
#     plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b')
    

plt.figure(1)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
    plt.plot(t,xx[:,i],'r',t,xx_cur[:,i],'b',t,xx1[:,i])
    # plt.plot(t,xx[:,i],'r',t,xx_cur[:,i],'b')
    plt.legend(['True value','estimation','Open loop'])

plt.figure(2)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
    plt.plot(t,a_est[:,i],'b',t,a[:,i],'r')
    # plt.plot(t,a_est[:,i],'b')
    plt.legend(['estimation','True'])
    
# plt.figure(4)
# for i in range (2):
#     plt.subplot(2,2,i+1)
#     plt.plot(t,a_est[:,i+14],'b',t,a[:,i+14],'r')
#     plt.legend(['estimation','True value'])
plt.show()



np.savetxt('xx_EM.txt',xx_cur)
np.savetxt('a_EM.txt',a)
np.savetxt('aest_EM.txt',a_est)
np.savetxt('openloop_EM.txt',xx1)


# np.savetxt('xx_EKF.txt',xx_cur)
# np.savetxt('a_EKF.txt',a)
# np.savetxt('aest_EKF.txt',a_est)
# np.savetxt('openloop_EKF.txt',xx1)